import sys
import paho.mqtt.client as mqtt
import tkinter as tk
from tkinter import ttk
from PIL import Image, ImageTk
import cv2
import numpy as np
import threading
import io
import base64
import time

class MqttImageReceiver:
    def __init__(self, root):
        self.root = root
        self.root.title("MQTT图像接收器")
        self.root.geometry("800x800")
        
        # 设置主题样式
        style = ttk.Style()
        style.configure("TFrame", background="#f5f5f5")
        style.configure("TLabel", background="#f5f5f5", font=('微软雅黑', 10))
        style.configure("Connected.TLabel", foreground="#28a745", font=('微软雅黑', 10, 'bold'))
        style.configure("Disconnected.TLabel", foreground="#dc3545", font=('微软雅黑', 10, 'bold'))
        
        # 创建主框架
        self.main_frame = ttk.Frame(root, padding="15")
        self.main_frame.pack(fill=tk.BOTH, expand=True)
        
        # 添加消息显示框
        self.message_frame = ttk.Frame(self.main_frame)
        self.message_frame.pack(fill=tk.BOTH, expand=True, pady=5)
        
        # 创建文本框和滚动条
        self.message_text = tk.Text(self.message_frame, height=6, width=50)
        self.scrollbar = ttk.Scrollbar(self.message_frame, orient="vertical", command=self.message_text.yview)
        self.message_text.configure(yscrollcommand=self.scrollbar.set)
        
        # 放置文本框和滚动条
        self.message_text.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        self.scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        
        # 设置文本框只读
        self.message_text.configure(state='disabled')
        
        # 添加主题输入框
        self.topic_frame = ttk.Frame(self.main_frame)
        self.topic_frame.pack(pady=5)
        
        ttk.Label(self.topic_frame, text="订阅主题:").pack(side=tk.LEFT, padx=5)
        self.topic_var = tk.StringVar(value="/image")
        self.topic_entry = ttk.Entry(self.topic_frame, textvariable=self.topic_var, width=30)
        self.topic_entry.pack(side=tk.LEFT, padx=5)
        
        self.subscribe_button = ttk.Button(self.topic_frame, text="订阅", command=self.subscribe_topic)
        self.subscribe_button.pack(side=tk.LEFT, padx=5)
        
        # 添加Broker IP输入框
        self.broker_frame = ttk.Frame(self.main_frame)
        self.broker_frame.pack(pady=5)
        
        ttk.Label(self.broker_frame, text="发送端 IP:").pack(side=tk.LEFT, padx=5)
        self.broker_ip_var = tk.StringVar(value="192.168.18.18")  # 修改默认IP地址
        self.broker_ip_entry = ttk.Entry(self.broker_frame, textvariable=self.broker_ip_var, width=20)
        self.broker_ip_entry.pack(side=tk.LEFT, padx=5)
        
        # 创建状态标签
        self.status_label = ttk.Label(
            self.main_frame, 
            text="等待接收图像...",
            font=('Arial', 12, 'bold')
        )
        self.status_label.pack(pady=5)
        
        # 创建图像显示区域
        self.image_label = ttk.Label(self.main_frame)
        self.image_label.pack(pady=10)
        
        # 创建连接状态指示器
        self.connection_frame = ttk.Frame(self.main_frame)
        self.connection_frame.pack(pady=5)
        
        self.connection_indicator = ttk.Label(
            self.connection_frame,
            text="●",
            style="Disconnected.TLabel",
            font=('Arial', 14)
        )
        self.connection_indicator.pack(side=tk.LEFT, padx=5)
        
        self.connection_status = ttk.Label(
            self.connection_frame,
            text="未连接",
            style="Disconnected.TLabel"
        )
        self.connection_status.pack(side=tk.LEFT)
        
        # 初始化MQTT客户端
        self.setup_mqtt()
        
        # 添加计时器变量
        self.last_image_time = 0
        self.check_timer = None
        self.start_image_check()
        
    def setup_mqtt(self):
        self.client = mqtt.Client()
        self.client.on_connect = self.on_connect
        self.client.on_message = self.on_message
        self.client.on_disconnect = self.on_disconnect
        
        # 添加重连标志和计数器
        self.is_connected = False
        self.reconnect_count = 0
        self.max_reconnect_attempts = 5
        
        # 在新线程中启动MQTT循环
        mqtt_thread = threading.Thread(target=self.mqtt_loop)
        mqtt_thread.daemon = True
        mqtt_thread.start()
        
    def mqtt_loop(self):
        while self.reconnect_count < self.max_reconnect_attempts:
            try:
                if not self.is_connected:
                    broker_ip = self.broker_ip_var.get()  # 获取用户输入的IP地址
                    print(f"尝试连接MQTT服务器 {broker_ip}... (尝试 {self.reconnect_count + 1}/{self.max_reconnect_attempts})")
                    self.client.connect(broker_ip, 1883, 60)
                    self.is_connected = True
                    self.reconnect_count = 0  # 重置重连计数
                    self.client.loop_forever()
            except Exception as e:
                print(f"MQTT连接错误: {e}")
                self.is_connected = False
                self.reconnect_count += 1
                self.root.after(0, self.update_status_label, f"连接失败，{self.max_reconnect_attempts - self.reconnect_count}次重试机会remaining...")
                time.sleep(5)  # 等待5秒后重试
        
        if self.reconnect_count >= self.max_reconnect_attempts:
            self.root.after(0, self.update_status_label, "无法连接到MQTT服务器，请检查服务器状态后重启应用")
        
    def on_disconnect(self, client, userdata, rc):
        print(f"与MQTT服务器断开连接，状态码: {rc}")
        self.is_connected = False
        
        # 更新连接状态显示
        self.root.after(0, self.update_connection_status, False)
        self.root.after(0, self.status_label.configure, {"text": "已断开连接，尝试重新连接..."})
        
        if rc != 0:
            try:
                self.client.reconnect()
            except Exception as e:
                print(f"重连失败: {e}")

    def subscribe_topic(self):
        """订阅新主题并重新连接到新的broker"""
        new_topic = self.topic_var.get()
        new_broker_ip = self.broker_ip_var.get()
        
        if new_topic and new_broker_ip:
            # 断开现有连接
            if self.is_connected:
                self.client.disconnect()
                self.is_connected = False
            
            # 重置重连计数器
            self.reconnect_count = 0
            
            # 更新状态
            self.status_label.configure(text="正在重新连接到新的服务器...")
            
            # 启动新的连接
            try:
                self.client.connect(new_broker_ip, 1883, 60)
                self.client.subscribe(new_topic)
                self.status_label.configure(text=f"已连接到 {new_broker_ip} 并订阅主题: {new_topic}")
            except Exception as e:
                self.status_label.configure(text=f"连接失败: {str(e)}")
                print(f"连接错误: {e}")

    def on_connect(self, client, userdata, flags, rc):
        print(f"已连接到MQTT服务器，状态码: {rc}")
        self.is_connected = True
        topic = self.topic_var.get()
        self.client.subscribe(topic)
        
        # 更新连接状态显示
        self.root.after(0, self.update_connection_status, True)
        self.root.after(0, self.status_label.configure, {"text": f"已连接到MQTT服务器，订阅主题: {topic}"})

    def on_message(self, client, userdata, msg):
        try:
            # 更新消息显示
            self.root.after(0, self.update_message_display, f"收到消息 - Topic: {msg.topic}\n载荷长度: {len(msg.payload)} 字节\n")
            
            print(f"收到消息，topic: {msg.topic}")
            print(f"收到消息，payload长度: {len(msg.payload)}")
            
            # 首先将消息解码为字符串，然后进行base64解码
            img_data = base64.b64decode(msg.payload)
            
            # 将base64解码后的数据转换为numpy数组
            nparr = np.frombuffer(img_data, np.uint8)
            
            # 调试信息
            print(f"base64解码后数据长度: {len(img_data)}")
            print(f"numpy数组形状: {nparr.shape}")
            
            # 尝试保存原始数据用于调试
            with open("debug_image.jpg", "wb") as f:
                f.write(img_data)
            
            # 使用OpenCV解码JPEG数据
            img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
            if img is None:
                raise ValueError("OpenCV无法解码图像数据")
            
            # BGR转RGB
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            
            # 转换为PIL图像
            pil_img = Image.fromarray(img)
            
            # 计算缩放比例
            display_size = (640, 480)
            pil_img.thumbnail(display_size, Image.Resampling.LANCZOS)
            
            # 更新最后接收图像的时间
            self.last_image_time = time.time()
            
            # 更新界面
            photo = ImageTk.PhotoImage(pil_img)
            self.root.after(0, self.update_ui, photo, msg.topic)
            
        except Exception as e:
            print(f"图像处理错误: {str(e)}")
            print(f"错误类型: {type(e)}")
            import traceback
            traceback.print_exc()
            
    def update_ui(self, photo, topic):
        self.image_label.configure(image=photo)
        self.image_label.image = photo  # 保持引用
        self.status_label.configure(text=f"当前Topic: {topic}")
    
    def start_image_check(self):
        """开始检查图像更新定时器"""
        self.check_timer = self.root.after(1000, self.check_image_timeout)
    
    def check_image_timeout(self):
        """检查图像是否超时"""
        current_time = time.time()
        if self.last_image_time > 0 and (current_time - self.last_image_time) > 2:
            self.clear_image()
        self.check_timer = self.root.after(1000, self.check_image_timeout)
    
    def clear_image(self):
        """清空图像显示"""
        self.image_label.configure(image='')
        self.image_label.image = None
        self.status_label.configure(text="无图像，等待重新连接...")
        self.last_image_time = 0
    
    def update_status_label(self, text):
        """更新状态标签的辅助方法"""
        self.status_label.configure(text=text)
    
    def update_message_display(self, message):
        """更新消息显示框"""
        self.message_text.configure(state='normal')
        # 添加分隔线
        separator = "-" * 50 + "\n"
        timestamp = time.strftime('%Y-%m-%d %H:%M:%S')
        formatted_message = f"{separator}[{timestamp}]\n{message}\n"
        self.message_text.insert(tk.END, formatted_message)
        self.message_text.see(tk.END)
        self.message_text.configure(state='disabled')

    def update_connection_status(self, is_connected):
        """更新连接状态指示器"""
        if is_connected:
            self.connection_indicator.configure(
                style="Connected.TLabel",
                text="● 已连接"
            )
            self.connection_status.configure(
                style="Connected.TLabel",
                text=f"MQTT服务器: {self.broker_ip_var.get()}:{1883}"
            )
        else:
            self.connection_indicator.configure(
                style="Disconnected.TLabel",
                text="● 未连接"
            )
            self.connection_status.configure(
                style="Disconnected.TLabel",
                text="MQTT服务器连接已断开"
            )

if __name__ == '__main__':
    root = tk.Tk()
    app = MqttImageReceiver(root)
    root.mainloop()
